# coding:utf-8
# author: 湖南涉外经济学院信息与机电工程学院
# date: 2022.09.09
from werkzeug.utils import secure_filename
import os
import time
from flask import Blueprint, render_template, request, send_file, redirect
import requests
import json
import cv2
import base64
import numpy as np
import zipfile
index_koutu = Blueprint("koutu", __name__)
# 设置允许的文件格式
ALLOWED_EXTENSIONS = set(['png', 'jpg', 'bmp', 'jpeg'])
# 获取当前运行脚本的绝对路径
basepath = os.path.dirname(__file__)


def cv2_to_base64(image):
    data = cv2.imencode('.jpg', image)[1]  # 将图片进行编码，并保存到本地
    return base64.b64encode(data.tobytes()).decode('utf8')


def base64_to_cv2(b64str):
    data = base64.b64decode(b64str.encode('utf8'))
    data = np.frombuffer(data, np.uint8)
    data = cv2.imdecode(data, cv2.IMREAD_COLOR)  # 读取网络url图像数据，并转换成图片格式：
    return data


def allowed_file(filename):
    filename = filename.lower()
    return '.' in filename and filename.rsplit('.', 1)[1] in ALLOWED_EXTENSIONS


def upload(f, save_path):
    print('---------欢迎使用抠图服务,图像分割模型为deeplabv3p_xception65_humanseg--------')
    print('step4 图片名:', f.filename)
    if not (f and allowed_file(f.filename)):
        # return jsonify({"error": 1001, "msg": "请检查上传的图片类型，仅限于png、PNG、jpg、JPG、bmp"})
        return render_template('404.html')
    sourcefile = os.path.join('static/images/source/',
                              secure_filename(f.filename))  # 将会对传入的文件路径进行处理，将其中的路径符号“/”用下划线代替，防止被黑客恶意攻击
    print('step5 上传图片保存在:', sourcefile)
    upload_path = os.path.join(basepath, sourcefile)  # 注意：没有的文件夹一定要先创建，不然会提示没有该路径
    f.save(upload_path)
    filename = convert(upload_path, f, save_path)
    print('step6 抠图完成,保存在:', filename)
    return sourcefile, filename, upload_path


# 压缩文件夹
def do_zip_compress(dirpath):
    output_name = f"{dirpath}.zip"
    parent_name = os.path.dirname(dirpath)

    zip = zipfile.ZipFile(output_name, "w", zipfile.ZIP_DEFLATED)
    # 多层级压缩
    for root, dirs, files in os.walk(dirpath):
        for file in files:
            if str(file).startswith("~$"):
                continue
            filepath = os.path.join(root, file)

            writepath = os.path.relpath(filepath, parent_name)
            zip.write(filepath, writepath)


# 去除背景色-
def convert(upload_path, f, save_path):
    # 发送HTTP请求
    data = {'images': [cv2_to_base64(cv2.imread(upload_path))]}
    headers = {"Content-type": "application/json"}
    # 指定图片分割方法为deeplabv3p_xception65_humanseg并发送post请求
    url = "http://127.0.0.1:8866/predict/deeplabv3p_xception65_humanseg"
    r = requests.post(url=url, headers=headers, data=json.dumps(data))  # 保存图片
    org_im = cv2.imread(upload_path)
    if r.json()["msg"]:  # 判断使用的CPU还是GPU
        mask = cv2.cvtColor(base64_to_cv2(r.json()["msg"][0]['data']), cv2.COLOR_BGR2GRAY)  # GPU
        print('step3 使用GPU训练')
    else:
        mask = cv2.cvtColor(base64_to_cv2(r.json()["results"][0]['data']), cv2.COLOR_BGR2GRAY)  # CPU
        print('step3 使用CPU训练')
    rgba = np.concatenate((org_im, np.expand_dims(mask, axis=2)), axis=2)
    filename = f.filename
    if save_path == None:
        abs_filename = os.path.join('static/images/target/', f.filename)
    else:
        abs_filename = os.path.join('static/images/target/', filename.split("/")[0], filename.split("/")[1])
    # abs_filename = os.path.join('static/images/target', secure_filename(''.join(lazy_pinyin(filename))))
    cv2.imencode('.png', rgba)[1].tofile(abs_filename)
    return abs_filename


# 上传并抠图
@index_koutu.route('/koutu_p', methods=['POST', 'GET'])  # 添加路由
def index_p():
    global save_path2
    if request.method == 'POST':
        try:
            time_start = time.time()
            f = request.files.getlist('file')
            save_path2 = f[0].filename.split("/")[0]
            save_path = os.path.join('static/images/target/', save_path2)
            save_path1 = save_path
            if os.path.exists(save_path):
                save_path = 'save_path'
            else:
                os.mkdir(save_path)
            for i in f:
                upload(i, save_path)
            time_end = time.time()
            time_c = time_end - time_start
            print('step5 所用时间: %.4f' % time_c, 's')
            do_zip_compress(save_path1)
            return redirect('/download/' + save_path2)

            # return render_template('koutu.html', wenjinmin=save_path2)
        except Exception as e:
            print(e)
            return render_template('404.html')
    return render_template('koutu.html')


@index_koutu.route('/koutu', methods=['POST', 'GET'])  # 添加路由
def index():
    if request.method == 'POST':
        try:
            time_start = time.time()
            f = request.files['file2']
            save_path = None
            sourcefile, filename, upload_path = upload(f, save_path)
            time_end = time.time()
            time_c = time_end - time_start
            print('step9 所用时间: %.4f' % time_c, 's')
            return render_template('koutu.html', sourcefile=sourcefile, filename=filename)
        except Exception as e:
            print(e)
            return render_template('404.html')
    return render_template('koutu.html')


@index_koutu.route('/download/<filename>', methods=['GET'])
def recive_file(filename):
    if request.method == "GET":
        try:
            return send_file("static/images/target/" + filename + ".zip", as_attachment=True)
        except Exception as e:
            print(e)
            return render_template('404.html')
